#!/usr/bin/env python
# -*- coding: utf-8 -*-

import asyncio
import subprocess
from pathlib import Path
from typing import List, Optional, Tuple, Dict, Any
import logging
import os
import uuid
import json
import copy

from .base_converter import BaseConverter, FileInfo

logger = logging.getLogger(__name__)

class PDFConverter(BaseConverter):
    """PDF转换器"""
    
    # 必需依赖应该为空
    REQUIRED_DEPENDENCIES = {}
    
    # 可选依赖定义
    OPTIONAL_DEPENDENCIES = {
        'marker': 'marker',
        'mineru': 'mineru',
        'docling': 'docling'
    }
    
    def __init__(self, temp_dir: Path, file_id: Optional[str] = None):
        super().__init__(temp_dir, file_id)
        self.default_method = self._get_default_method()
    
    def _get_default_method(self) -> str:
        """获取默认转换方法"""
        # 优先级：marker > docling > mineru
        for method in ['marker', 'docling', 'mineru']:
            if self.is_dependency_available(method):
                return method
        
        # 如果都不可用，返回marker作为默认值（会在实际使用时报错）
        return 'marker'
    
    async def convert(self, file_path: Path, method: Optional[str] = None, **kwargs) -> Tuple[str, List[FileInfo]]:
        """
        转换PDF文件
        
        Args:
            file_path: PDF文件路径
            method: 转换方法 ('marker', 'mineru', 'docling')
            **kwargs: 其他参数
            
        Returns:
            tuple: (markdown内容, 文件信息列表)
        """
        method = method or self.default_method
        
        logger.info(f"使用 {method} 转换PDF: {file_path.name}")
        
        # 检查方法是否可用
        if not self.is_dependency_available(method):
            available_methods = self.get_available_methods()
            if available_methods:
                fallback_method = available_methods[0]
                logger.warning(f"{method} 不可用，使用 {fallback_method} 作为替代")
                method = fallback_method
            else:
                raise ImportError(
                    f"没有可用的PDF转换器。请安装以下任一依赖：\n"
                    f"- marker-pdf (推荐): pip install marker-pdf\n"
                    f"- mineru: 需要特殊安装方式，先安装uv然后运行 'uv pip install -U \"mineru[core]\"'\n"
                    f"- docling: pip install docling"
                )
        
        if method == 'marker':
            return await self._convert_with_marker(file_path, **kwargs)
        elif method == 'mineru':
            return await self._convert_with_mineru(file_path, **kwargs)
        elif method == 'docling':
            return await self._convert_with_docling(file_path, **kwargs)
        else:
            raise ValueError(f"不支持的转换方法: {method}")
    
    async def _convert_with_marker(self, file_path: Path, **kwargs) -> Tuple[str, List[FileInfo]]:
        """使用marker转换PDF - 使用新版Python API"""
        self.require_dependency('marker', 'marker-pdf')
        
        try:
            # 使用新版Marker API
            try:
                from marker.converters.pdf import PdfConverter  # type: ignore
                from marker.models import create_model_dict  # type: ignore
                from marker.output import text_from_rendered  # type: ignore
                
                logger.info("使用新版marker Python API转换PDF")
                
                # 创建模型字典
                model_dict = create_model_dict()
                
                # 创建转换器
                converter = PdfConverter(
                    artifact_dict=model_dict,
                )
                
                # 执行转换
                rendered = converter(str(file_path))
                
                # 提取文本和图片
                text, _, images = text_from_rendered(rendered)
                
                # 创建输出目录
                file_stem = file_path.stem
                output_dir = self.converter_work_dir / f"marker_{file_stem}"
                images_dir = output_dir / "images"
                images_dir.mkdir(parents=True, exist_ok=True)
                
                # 生成markdown内容
                markdown_content = f"# PDF文档转换结果\n\n"
                markdown_content += f"**文件名**: {file_path.name}\n\n"
                markdown_content += f"**转换工具**: Marker (新版Python API)\n\n"
                markdown_content += f"## 文档内容\n\n{text}"
                
                # 处理图片和文件信息
                file_infos = []
                
                # 保存图片到本地目录
                if images:
                    logger.info(f"保存 {len(images)} 张图片到本地目录")
                    for img_name, img_data in images.items():
                        img_path = images_dir / img_name
                        # 使用基类的通用图片保存方法
                        if await self.save_image_data(img_data, img_path):
                            img_file_info = self.create_file_info(img_path, str(uuid.uuid4()), "image")
                            file_infos.append(img_file_info)
                
                # 保存markdown文件
                md_path = output_dir / f"{file_stem}.md"
                with open(md_path, 'w', encoding='utf-8') as f:
                    f.write(markdown_content)
                
                # 添加文档信息
                doc_file_info = self.create_file_info(md_path, str(uuid.uuid4()), "document")
                file_infos.append(doc_file_info)
                
                logger.info(f"marker转换成功: {len(file_infos)} 个文件")
                return markdown_content, file_infos
                
            except ImportError as ie:
                logger.warning(f"新版marker API不可用: {ie}，尝试旧版API")
                return await self._marker_legacy_api(file_path, **kwargs)
            
        except Exception as e:
            logger.error(f"marker转换失败: {e}")
            raise RuntimeError(f"marker转换失败: {str(e)}")
    
    async def _marker_legacy_api(self, file_path: Path, **kwargs) -> Tuple[str, List[FileInfo]]:
        """使用旧版marker API作为回退"""
        try:
            from marker.convert import convert_single_pdf
            from marker.models import load_all_models
            
            logger.info("使用旧版marker Python API转换PDF")
            
            # 加载模型
            model_list = load_all_models()
            
            # 执行转换
            full_text, images, out_meta = convert_single_pdf(
                str(file_path), 
                model_list, 
                max_pages=kwargs.get('max_pages', None),
                langs=kwargs.get('langs', None)
            )
            
            # 创建输出目录
            file_stem = file_path.stem
            output_dir = self.converter_work_dir / f"marker_{file_stem}"
            images_dir = output_dir / "images"
            images_dir.mkdir(parents=True, exist_ok=True)
            
            # 生成markdown内容
            markdown_content = f"# PDF文档转换结果\n\n"
            markdown_content += f"**文件名**: {file_path.name}\n\n"
            markdown_content += f"**转换工具**: Marker (旧版Python API)\n\n"
            if out_meta:
                markdown_content += f"**页数**: {out_meta.get('total_pages', 'N/A')}\n\n"
            markdown_content += f"## 文档内容\n\n{full_text}"
            
            # 处理图片和文件信息
            file_infos = []
            
            # 保存图片到本地目录
            if images:
                logger.info(f"保存 {len(images)} 张图片到本地目录")
                for img_name, img_data in images.items():
                    img_path = images_dir / img_name
                    # 使用基类的通用图片保存方法
                    if await self.save_image_data(img_data, img_path):
                        img_file_info = self.create_file_info(img_path, str(uuid.uuid4()), "image")
                        file_infos.append(img_file_info)
            
            # 保存markdown文件
            md_path = output_dir / f"{file_stem}.md"
            with open(md_path, 'w', encoding='utf-8') as f:
                f.write(markdown_content)
            
            # 添加文档信息
            doc_file_info = self.create_file_info(md_path, str(uuid.uuid4()), "document")
            file_infos.append(doc_file_info)
            
            logger.info(f"marker转换成功: {len(file_infos)} 个文件")
            return markdown_content, file_infos
            
        except ImportError:
            logger.warning("旧版marker API也不可用，尝试命令行方式")
            return await self._marker_command_fallback(file_path, **kwargs)
        except Exception as e:
            logger.error(f"旧版marker API转换失败: {e}")
            return await self._marker_command_fallback(file_path, **kwargs)
    
    async def _marker_command_fallback(self, file_path: Path, **kwargs) -> Tuple[str, List[FileInfo]]:
        """marker命令行回退方案"""
        try:
            file_stem = file_path.stem
            output_dir = self.converter_work_dir / f"marker_{file_stem}"
            output_dir.mkdir(exist_ok=True)
            
            # 构建marker命令
            cmd = [
                'marker_single', 
                str(file_path),
                str(output_dir),
                '--batch_multiplier', '2',
                '--max_pages', str(kwargs.get('max_pages', 200)),
                '--langs', kwargs.get('langs', 'Chinese,English')
            ]
            
            logger.info(f"执行marker命令: {' '.join(cmd)}")
            
            # 执行转换命令
            process = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                cwd=str(output_dir)
            )
            
            stdout, stderr = await process.communicate()
            
            if process.returncode != 0:
                logger.error(f"marker命令执行失败: {stderr.decode()}")
                return await self._marker_basic_fallback(file_path, **kwargs)
            
            # 查找生成的markdown文件
            md_files = list(output_dir.glob("*.md"))
            if not md_files:
                return await self._marker_basic_fallback(file_path, **kwargs)
            
            # 读取markdown内容
            md_file = md_files[0]
            markdown_content = md_file.read_text(encoding='utf-8')
            
            # 查找生成的图片文件
            file_infos = []
            
            # 创建主文档的文件信息
            doc_file_info = self.create_file_info(md_file, str(uuid.uuid4()), "document")
            file_infos.append(doc_file_info)
            
            # 查找图片文件
            for img_pattern in ['*.png', '*.jpg', '*.jpeg', '*.gif']:
                for img_file in output_dir.rglob(img_pattern):
                    img_file_info = self.create_file_info(img_file, str(uuid.uuid4()), "image")
                    file_infos.append(img_file_info)
            
            logger.info(f"marker转换成功: {len(file_infos)} 个文件")
            return markdown_content, file_infos
            
        except Exception as e:
            logger.warning(f"marker命令行执行失败: {e}")
            return await self._marker_basic_fallback(file_path, **kwargs)
    
    async def _marker_basic_fallback(self, file_path: Path, **kwargs) -> Tuple[str, List[FileInfo]]:
        """marker基本回退方案"""
        markdown_content = f"# PDF文档\n\n**文件名**: {file_path.name}\n\n**说明**: marker转换器需要安装完整的依赖。请运行: pip install marker-pdf"
        
        file_infos = []
        doc_file_info = self.create_file_info(file_path, str(uuid.uuid4()), "document")
        file_infos.append(doc_file_info)
        
        return markdown_content, file_infos
    
    async def _convert_with_mineru(self, file_path: Path, **kwargs) -> Tuple[str, List[FileInfo]]:
        """使用mineru转换PDF - 使用正确的Python API"""
        self.require_dependency('mineru', 'mineru')
        
        try:
            # 使用正确的MinerU API
            try:
                from mineru.cli.common import convert_pdf_bytes_to_bytes_by_pypdfium2, prepare_env, read_fn
                from mineru.data.data_reader_writer import FileBasedDataWriter
                from mineru.utils.enum_class import MakeMode
                from mineru.backend.pipeline.pipeline_analyze import doc_analyze as pipeline_doc_analyze
                from mineru.backend.pipeline.pipeline_middle_json_mkcontent import union_make as pipeline_union_make
                from mineru.backend.pipeline.model_json_to_middle_json import result_to_middle_json as pipeline_result_to_middle_json
                
                logger.info("使用正确的mineru Python API转换PDF")
                
                # 创建输出目录
                file_stem = file_path.stem
                output_base_dir = self.converter_work_dir / f"mineru_{file_stem}"
                output_base_dir.mkdir(exist_ok=True)
                
                # 读取PDF文件
                pdf_bytes = read_fn(file_path)
                pdf_bytes_list = [pdf_bytes]
                pdf_file_names = [file_stem]
                lang_list = [kwargs.get('lang', 'ch')]
                
                # 配置参数
                backend = kwargs.get('backend', 'pipeline')
                parse_method = kwargs.get('parse_method', 'auto')
                formula_enable = kwargs.get('formula_enable', True)
                table_enable = kwargs.get('table_enable', True)
                start_page_id = kwargs.get('start_page_id', 0)
                end_page_id = kwargs.get('end_page_id', None)
                
                # 处理PDF字节数据
                new_pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id, end_page_id)
                pdf_bytes_list[0] = new_pdf_bytes
                
                # 分析PDF
                infer_results, all_image_lists, all_pdf_docs, lang_list, ocr_enabled_list = pipeline_doc_analyze(
                    pdf_bytes_list, 
                    lang_list, 
                    parse_method=parse_method, 
                    formula_enable=formula_enable,
                    table_enable=table_enable
                )
                
                # 处理结果
                model_list = infer_results[0]
                images_list = all_image_lists[0]
                pdf_doc = all_pdf_docs[0]
                _lang = lang_list[0]
                _ocr_enable = ocr_enabled_list[0]
                
                # 准备环境
                local_image_dir, local_md_dir = prepare_env(str(output_base_dir), file_stem, parse_method)
                image_writer = FileBasedDataWriter(local_image_dir)
                md_writer = FileBasedDataWriter(local_md_dir)
                
                # 转换为中间JSON格式
                middle_json = pipeline_result_to_middle_json(
                    model_list, images_list, pdf_doc, image_writer, _lang, _ocr_enable, formula_enable
                )
                
                pdf_info = middle_json["pdf_info"]
                
                # 生成markdown内容
                image_dir_name = os.path.basename(local_image_dir)
                md_content_str = pipeline_union_make(pdf_info, MakeMode.MM_MD, image_dir_name)
                
                # 保存markdown文件
                md_filename = f"{file_stem}.md"
                md_writer.write_string(md_filename, md_content_str)
                
                # 检查生成的图片文件，如果markdown中没有图片引用，就手动添加
                image_dir_path = Path(local_image_dir)
                generated_images = []
                if image_dir_path.exists():
                    for img_file in image_dir_path.glob("*.jpg"):
                        generated_images.append(img_file.name)
                    for img_file in image_dir_path.glob("*.png"):
                        generated_images.append(img_file.name)
                
                # 如果有图片但markdown中没有引用，则添加引用
                import re
                existing_refs = re.findall(r'!\[.*?\]\([^)]+\)', md_content_str)
                if generated_images and not existing_refs:
                    logger.info(f"mineru生成了 {len(generated_images)} 个图片，但markdown中没有引用，手动添加")
                    image_section = "\n\n## 提取的图片\n\n"
                    for i, img_name in enumerate(generated_images, 1):
                        img_relative_path = f"{image_dir_name}/{img_name}"
                        image_section += f"![图片 {i}]({img_relative_path})\n\n"
                    md_content_str += image_section
                
                # 添加头部信息到markdown内容
                markdown_content = f"# PDF文档转换结果\n\n"
                markdown_content += f"**文件名**: {file_path.name}\n\n"
                markdown_content += f"**转换工具**: MinerU (Python API)\n\n"
                markdown_content += f"## 文档内容\n\n{md_content_str}"
                
                # 收集文件信息
                file_infos = []
                
                # 添加主markdown文件
                md_file_path = Path(local_md_dir) / md_filename
                if md_file_path.exists():
                    doc_file_info = self.create_file_info(md_file_path, str(uuid.uuid4()), "document")
                    file_infos.append(doc_file_info)
                
                # 使用基类方法收集生成的图片文件
                image_dir_path = Path(local_image_dir)
                if image_dir_path.exists():
                    generated_files = self.collect_generated_files(image_dir_path, ['*.png', '*.jpg', '*.jpeg', '*.gif'])
                    file_infos.extend(generated_files)
                
                # 保存其他有用的文件
                try:
                    # 保存content list
                    content_list = pipeline_union_make(pdf_info, MakeMode.CONTENT_LIST, image_dir_name)
                    md_writer.write_string(
                        f"{file_stem}_content_list.json",
                        json.dumps(content_list, ensure_ascii=False, indent=4),
                    )
                    
                    # 保存middle json
                    md_writer.write_string(
                        f"{file_stem}_middle.json",
                        json.dumps(middle_json, ensure_ascii=False, indent=4),
                    )
                except Exception as e:
                    logger.warning(f"保存额外文件失败: {e}")
                
                logger.info(f"mineru转换成功: {len(file_infos)} 个文件")
                return markdown_content, file_infos
                
            except ImportError as ie:
                logger.warning(f"mineru Python API不可用: {ie}，尝试命令行方式")
                return await self._mineru_command_fallback(file_path, **kwargs)
            
        except Exception as e:
            logger.error(f"mineru转换失败: {e}")
            return await self._mineru_fallback(file_path, **kwargs)
    
    async def _mineru_command_fallback(self, file_path: Path, **kwargs) -> Tuple[str, List[FileInfo]]:
        """mineru命令行回退方案"""
        try:
            file_stem = file_path.stem
            output_dir = self.converter_work_dir / f"mineru_{file_stem}"
            output_dir.mkdir(exist_ok=True)
            
            # 使用处理.txt中验证过的mineru命令
            cmd = [
                'magic-pdf',
                'pdf-command',
                '--pdf', str(file_path),
                '--output', str(output_dir)
            ]
            
            logger.info(f"执行mineru命令: {' '.join(cmd)}")
            
            process = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            
            stdout, stderr = await process.communicate()
            
            if process.returncode != 0:
                logger.warning(f"mineru命令执行失败: {stderr.decode()}")
                return await self._mineru_fallback(file_path, **kwargs)
            
            # 查找生成的markdown文件
            md_files = list(output_dir.rglob("*.md"))
            if md_files:
                md_file = md_files[0]
                markdown_content = md_file.read_text(encoding='utf-8')
            else:
                markdown_content = f"# PDF文档转换结果\n\n**文件名**: {file_path.name}\n\n**转换工具**: MinerU\n\n注意：未找到生成的markdown文件。"
            
            # 收集文件信息
            file_infos = []
            
            # 使用基类方法收集所有生成的文件
            generated_files = self.collect_generated_files(output_dir, ['*.md', '*.png', '*.jpg', '*.jpeg'])
            file_infos.extend(generated_files)
            
            if not file_infos:
                # 如果没有生成文件，至少添加原文档信息
                doc_file_info = self.create_file_info(file_path, str(uuid.uuid4()), "document")
                file_infos.append(doc_file_info)
            
            return markdown_content, file_infos
            
        except Exception as e:
            logger.warning(f"mineru命令行执行失败: {e}")
            return await self._mineru_fallback(file_path, **kwargs)
    
    async def _mineru_fallback(self, file_path: Path, **kwargs) -> Tuple[str, List[FileInfo]]:
        """mineru回退方案"""
        markdown_content = f"# PDF文档\n\n**文件名**: {file_path.name}\n\n**说明**: MinerU转换器需要正确的环境配置。请确保已正确安装mineru依赖。"
        
        file_infos = []
        doc_file_info = self.create_file_info(file_path, str(uuid.uuid4()), "document")
        file_infos.append(doc_file_info)
        
        return markdown_content, file_infos
    
    async def _convert_with_docling(self, file_path: Path, **kwargs) -> Tuple[str, List[FileInfo]]:
        """使用docling转换PDF文件"""
        logger.info(f"使用docling转换PDF文件: {file_path.name}")
        
        # 直接使用基类的通用docling方法
        return await self.convert_with_docling(file_path, **kwargs)
    
    def get_available_methods(self) -> List[str]:
        """获取可用的转换方法"""
        methods = []
        for method in ['marker', 'mineru', 'docling']:
            if self.is_dependency_available(method):
                methods.append(method)
        return methods